package org.kth.dks.util;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

/**
 * AsyncOperation implements a class to register operations,
 * to associate a state with an operation, and to wait for the completion
 * of operations. Multiple threads waiting for an operation to complete is possible.
 */
public class AsyncOperation {
    static Map map = Collections.synchronizedMap(new HashMap());
    static long keyCounter = 0;

    Future future;
    String key;
    Object state;

    AsyncOperation() { }

    /* Starts a new operation, creates a mapping from a name to
     * it, and returns the compound operation object.
     */
    public static AsyncOperation start() {
	return start(null);
    }

    /* Starts a new operation, creates a mapping from a name to
     * it, remembers a state associated with the operation,
     * and returns the compound operation object.
     */
    public static AsyncOperation start(Object state) {
	AsyncOperation op = new AsyncOperation();
	op.key = Long.toString(keyCounter++);
	op.future = new Future();
	op.state = state;
	map.put(op.key, op);
	return op;
    }

    /* Retrieves the name of an operation.
     */
    public String getKey() {
	return key;
    }

    /* Retrieves the state of an operation.
     */
    public Object getState() {
	return state;
    }

    /* Updates the state of an operation.
     */
    public void setState(Object state) {
	this.state = state;
    }

    /* Locates an operation based on its name.
     */
    public static AsyncOperation get(String key) {
	return (AsyncOperation) map.get(key);
    }

    /* Waits for an operation to complete; removes the corresponding
     * mapping on success.
     */
    public Object waitOn()
	throws InterruptedException, CancellationException, Exception {
	Object result = future.get();
	map.remove(key);
	return result;
    }

    /* Waits for an operation to complete, up to timeout milliseconds;
     * removes the corresponding
     * mapping on success or if the timeout expires.
     */
    public Object waitOn(long timeout)
	throws InterruptedException, TimeoutException, CancellationException {
	Object result = null;

	try { result = future.get(timeout); }
	catch (TimeoutException e) {
	    map.remove(key);
	    throw e;
	}

	map.remove(key);
	return result;
    }

    /* Complete an operation.
     */
    public void complete(Object result) {
	future.store(result);
    }

    /* Complete an operation.
     */
    public static void complete(String name, Object result) {
	AsyncOperation op = AsyncOperation.get(name);
	op.future.store(result);
    }

    /* Void an operation mapping. Removes the mapping.
     */
    public void cancel() {
	future.cancel(false);
	map.remove(key);
    }

    /* Void an operation mapping. Removes the mapping.
     */
    public void cancel(Exception ex) {
        future.cancel(ex);
        map.remove(key);
    }

    /* Void an operation mapping.
     */
    public void voidOperationMapping() {
	map.remove(key);
    }
}
